#include "./Base.h"
#include  <ntifs.h>


#ifndef _WIN64


//����KiProcessorBlock�����ַ
PVOID  SearchKiProcessorBlockAddress()
{
	UCHAR buf[] = {
	0x8b,0x04,0x85
	};
	PVOID addressB;
	PVOID KiProcessorBlock;
	PVOID funcAddress;
	UNICODE_STRING funName;
	RtlInitUnicodeString(&funName, L"NtQuerySystemInformation");
	funcAddress = (PVOID)((ULONG)MmGetSystemRoutineAddress(&funName) + 0x176);
	addressB = SearchPatterFuzzy((PCH)funcAddress, 0x400, (PCH)buf, sizeof(buf));

	KiProcessorBlock = (void*)(*(ULONG*)((PCH)addressB + sizeof(buf)));
	if (!KiProcessorBlock) return NULL;
	return KiProcessorBlock;
}


//����PspCidTable��ַ
PVOID SearchPspCidTable()
{

	UCHAR buf[] = {
	0xFF,0x35
	};
	PVOID addressB;
	PVOID PspCidTable;
	PVOID funcAddress;
	UNICODE_STRING funName;
	RtlInitUnicodeString(&funName, L"PsLookupProcessThreadByCid");
	funcAddress = (PVOID)((ULONG)MmGetSystemRoutineAddress(&funName));
	addressB = SearchPatterFuzzy((PCH)funcAddress, 0x400, (PCH)buf, sizeof(buf));
	PspCidTable = (void*)(*(ULONG*)((PCH)addressB + sizeof(buf)));
	if (PspCidTable) return PspCidTable;
	return NULL;
}

VOID PrintProcessInfoMation(PSYSTEM_PROCESSES process) {

	INT  i;
	LOGL();
	LOG("[PROCESS]\r\n");
	LOG("\t[NAME]:%ls\t[PID]:%d\t[HandleCount]:%d\t[ThreadCount]:%d\r\n", process->ProcessName.Buffer, process->ProcessId, process->HandleCount, process->ThreadCount);
	LOG("\t[THREAD]\r\n");

	for (i = 0; i < process->ThreadCount; i++) {
		SYSTEM_THREADS thread = process->Threads[i];
		LOG("\t\t\t[TID]:%d\t[PRIORITY]:%d\r\n", thread.ClientIs.UniqueThread, thread.Priority);
	}
}
//��ȡϵͳ������Ϣ
VOID QuerySystemProcessInfomation() {

	UINT32 bufSize = 0;
	NTSTATUS status = ZwQuerySystemInformation(
		SystemProcessInformation,
		&bufSize,
		bufSize,
		(PULONG)&bufSize
	);
	if (!bufSize) {
		LOG("ZwQuerySystemInformation ��ȡ����ʧ�ܣ�\r\n");
		return;
	}
	PSYSTEM_PROCESSES pbuf = (PSYSTEM_PROCESSES)ExAllocatePoolWithTag(NonPagedPool, bufSize * 2, 'yang');
	if (!pbuf) {
		LOG("ExAllocatePoolWithTag fail \r\n");
	}
	memset(pbuf, 0, bufSize * 2);
	status = ZwQuerySystemInformation(
		SystemProcessInformation,
		pbuf,
		bufSize * 2,
		(PULONG)&bufSize
	);
	if (!NT_SUCCESS(status)) {
		LOG("ZwQuerySystemInformation ��ȡ��Ϣʧ�ܣ�\r\n");
		return;
	}
	do {

		PrintProcessInfoMation(pbuf);
		pbuf = (PSYSTEM_PROCESSES)((ULONG)pbuf + pbuf->NextEntryDelta);

	} while (pbuf->NextEntryDelta != 0);


}
VOID PrintModuleInformation(SYSTEM_MODULE_ENTRY module) {
	INT  i;
	LOGL();
	LOG("[MODULE]\r\n");
	LOG("\t[IMAGE BASE]:%x\t[IMAGE SIZE]:%x\t[LoadCount]:%d\r\n", module.ImageBase, module.ImageSize, module.LoadCount);
	LOG("\t[PATH]:%s\r\n", module.FullPathName);
}

//��ȡϵͳģ����Ϣ
VOID QuerySystemModuleInfomation() {

	UINT32 bufSize = 0;
	INT i;
	NTSTATUS status = ZwQuerySystemInformation(
		SystemModuleInformation,
		&bufSize,
		bufSize,
		(PULONG)&bufSize
	);
	if (!bufSize) {
		LOG("ZwQuerySystemInformation ��ȡ����ʧ�ܣ�\r\n");
		return;
	}
	PSYSTEM_MODULE_INFORMATION pbuf = (PSYSTEM_MODULE_INFORMATION)ExAllocatePoolWithTag(NonPagedPool, bufSize * 2, 'yang');
	if (!pbuf) {
		LOG("ExAllocatePoolWithTag fail \r\n");
	}
	memset(pbuf, 0, bufSize * 2);
	status = ZwQuerySystemInformation(
		SystemModuleInformation,
		pbuf,
		bufSize * 2,
		(PULONG)&bufSize
	);
	if (!NT_SUCCESS(status)) {
		LOG("ZwQuerySystemInformation ��ȡ��Ϣʧ�ܣ�\r\n");
		ExFreePoolWithTag(pbuf, 'yang');
		return;
	}
	for (i = 0; i < pbuf->Count; i++) {
		PrintModuleInformation(pbuf->Module[i]);
	}
	ExFreePoolWithTag(pbuf, 'yang');
}


#else


BOOLEAN IsWin7() {

	RTL_OSVERSIONINFOW version = { 0 };
	RtlGetVersion(&version);
	return (version.dwBuildNumber == 7600 || version.dwBuildNumber == 7601);
}

VOID KernelSleep(ULONG64 ms,BOOLEAN isAlertable) {
	LARGE_INTEGER inTime;
	inTime.QuadPart = -ms * 10 * 1000;
	KeDelayExecutionThread(KernelMode, isAlertable, &inTime);
}


//����KiProcessorBlock�����ַ
PVOID  SearchKiProcessorBlockAddress()
{

	char codeBuf[] = "488D0D****";
	return (PVOID)findSystemOffsetByOffsetOfExprotFunction(L"MmAllocateNonCachedMemory", codeBuf, 0x0, 0x3);
}


//����PspCidTable��ַ
PVOID SearchPspCidTable()
{

	char codeBuf[] = "B2*E8****";
	PVOID tempAddress = (PVOID)findSystemOffsetByOffsetOfExprotFunction(L"PsLookupProcessByProcessId", codeBuf, 0x0, 0x3);
	if (!tempAddress) return 0;
	char codeBuf2[] = "488B05****";
	tempAddress = (PVOID)findSystemOffsetByOffsetAddress(tempAddress, codeBuf2, 0x0, 0x3);
	return tempAddress;

}

PVOID SearchPspNotifyEnableMask() {

	char codeBuf[] = "8B*****A8*";
	return (PVOID)findSystemOffsetByOffsetOfExprotFunction(L"PsSetLoadImageNotifyRoutineEx", codeBuf, 0x0, 0x2);
	
}


#endif // !_WIN64



